"""
  描述：点质量模型（用于直行的车辆）
"""

# 外部包
import numpy as np

# 内部包
from kalman_filter._01_kf_point_mass_model import PointMassModelKalmanFilter

class PointMassModel:
  def __init__(self, dt, y0, x0, v0, process_noise_std=[0.1, 0.2, 0.1], measurement_noise_std=[0.5, 0.2, 0.2]):
    self.dt = dt  # 时间步长
    self.process_noise_std = process_noise_std
    self.measurement_noise_std = measurement_noise_std
    
    # 初始状态
    self.y = y0  # 纵向位置y
    self.x = x0  # 横向位置x
    self.v = v0  # 纵向速度v

    # 测量值
    self.measure_y = y0
    self.measure_x = x0
    self.measure_v = v0

    # 滤波状态
    self.kf = PointMassModelKalmanFilter(dt, process_noise_std, measurement_noise_std)
    self.kf.reset(np.array([y0, v0, 0.0])) # 初始状态估计
    self.filter_y, self.filter_v, _ = self.kf.get_state()
    self.filter_x = x0  # 横向位置x不变
    P = self.kf.get_covariance()
    self.filter_covariances = [P[0,0], P[1,1], P[2,2]] # 位置、速度、加速度的协方差
    
  def _add_noise(self, data, stds_dev, noise_type="gaussian"):
    """添加噪声到数据"""
    stds_dev = np.array(stds_dev)
    if noise_type == "gaussian":
        # 高斯噪声
        noise = np.random.normal(0, stds_dev, data.shape)
    elif noise_type == "uniform":
        # 均匀噪声范围为 [-std, std]
        noise = np.random.uniform(-stds_dev, stds_dev, data.shape)
    return data + noise

  def update(self, a):
    """
    使用点质量模型更新车辆状态
    参数:
      a: 纵向加速度
    """
    # 更新速度和位置
    self.y += self.v * self.dt + 0.5 * a * (self.dt**2)
    self.v += a * self.dt

    # 生成测量值
    measurement_data = self._add_noise(np.array([self.y, self.v, a]), self.measurement_noise_std) # 真实值增加噪声，形成观测值
    self.measure_y, self.measure_v, _ = measurement_data

    # 生成滤波状态
    self.kf.filter_step(measurement_data, a) # 执行滤波步骤
    self.filter_y, self.filter_v, _ = self.kf.get_state()
    P = self.kf.get_covariance()
    self.filter_covariances = [P[0,0], P[1,1], P[2,2]] # 位置、速度、加速度的协方差

  def get_state(self):
    """
    获取当前状态
    """
    return self.y, self.x, self.v
  
  def get_measure_state(self):
    """
    获取测量状态
    """
    return self.measure_y, self.measure_x, self.measure_v
  
  def get_filter_state(self):
    """
    获取滤波状态
    """
    return self.filter_y, self.filter_x, self.filter_v, self.filter_covariances
  
  def calc_predict_state(self, a_arr):
    """
    计算预测行为对应的状态
    参数:
      a: 纵向加速度数组
    """
    states = []
    # !在计算预测状态时，应该基于滤波状态进行计算。因为真实状态车辆无法感知到，当前帧应用的也是滤波状态
    y_pred = self.filter_y
    v_pred = self.filter_v
    for a in a_arr:
      y_pred += v_pred * self.dt + 0.5 * a * (self.dt**2)
      v_pred += a * self.dt
      states.append((y_pred, self.filter_x, v_pred))
    return np.array(states)
    